{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a4d5d1bc-df65-47ae-8e2a-d1fe6218fb3e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:32:59.113960Z",
     "iopub.status.busy": "2022-05-31T02:32:59.113739Z",
     "iopub.status.idle": "2022-05-31T02:33:00.778490Z",
     "shell.execute_reply": "2022-05-31T02:33:00.777934Z",
     "shell.execute_reply.started": "2022-05-31T02:32:59.113900Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.55 s, sys: 97.3 ms, total: 1.65 s\n",
      "Wall time: 1.66 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import Dataset\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.utils.data as Data\n",
    "from torchvision import datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "029bc639-dace-4661-875b-ec0605f5f8e2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:33:00.779669Z",
     "iopub.status.busy": "2022-05-31T02:33:00.779385Z",
     "iopub.status.idle": "2022-05-31T02:33:00.785470Z",
     "shell.execute_reply": "2022-05-31T02:33:00.784914Z",
     "shell.execute_reply.started": "2022-05-31T02:33:00.779645Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GPU:False\n"
     ]
    }
   ],
   "source": [
    "class Parse:\n",
    "    def __init__(self):  #定义超参数\n",
    "        self.channels = 1\n",
    "        self.img_size = 28\n",
    "        self.latent_dim = 100\n",
    "        self.epochs = 200\n",
    "        self.batch_size = 256\n",
    "        self.lr = 0.0002\n",
    "        self.b = (0.5, 0.999)\n",
    "        self.cpu_nums = 8\n",
    "        self.classes = 10\n",
    "        self.want_cuda = \"cuda:0\"\n",
    "        \n",
    "        \n",
    "\n",
    "        # 下面是衍生的\n",
    "        self.img_shape = (self.channels, self.img_size, self.img_size)\n",
    "        self.cuda = torch.cuda.is_available()\n",
    "        self.device = self.want_cuda if self.cuda else 'cpu'\n",
    "        self.big_cuda = True  # 能不能把数据集全部放入GPU\n",
    "        print(f\"GPU:{self.cuda}\")\n",
    "opt = Parse()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9a9197d7-5606-4b26-804b-c4e64a5f1354",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T08:07:24.064644Z",
     "iopub.status.busy": "2022-05-31T08:07:24.063992Z",
     "iopub.status.idle": "2022-05-31T08:07:24.070846Z",
     "shell.execute_reply": "2022-05-31T08:07:24.070257Z",
     "shell.execute_reply.started": "2022-05-31T08:07:24.064615Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Tools:\n",
    "    # 装一些杂七杂八的工具\n",
    "    def block(in_feat, out_feat, normalize=True):\n",
    "        layers = [nn.Linear(in_feat, out_feat),nn.LeakyReLU(0.2, inplace=True)]\n",
    "        if normalize:\n",
    "            layers.insert(1, nn.BatchNorm1d(out_feat, 0.8))\n",
    "        return layers\n",
    "    \n",
    "    \n",
    "    def sample_image(n_row, batches_done, generator):\n",
    "\n",
    "        z = torch.randn(n_row ** 2, opt.latent_dim, device=opt.device)\n",
    "        labels = torch.tensor([num for _ in range(n_row) for num in range(n_row)], device=opt.device)\n",
    "        gen_imgs = generator(z, labels)\n",
    "        save_image(gen_imgs.data, \"images1/%d.png\" % batches_done, nrow=n_row, normalize=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c8d46030-6bfa-40f4-a98e-9208c741d7b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:33:00.793965Z",
     "iopub.status.busy": "2022-05-31T02:33:00.793715Z",
     "iopub.status.idle": "2022-05-31T02:33:00.800367Z",
     "shell.execute_reply": "2022-05-31T02:33:00.799816Z",
     "shell.execute_reply.started": "2022-05-31T02:33:00.793944Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.label_emb = nn.Embedding(opt.classes, opt.classes)\n",
    "        \n",
    "        self.model = nn.Sequential(\n",
    "            *Tools.block(opt.latent_dim+opt.classes, 128, normalize=False),\n",
    "            *Tools.block(128, 256),\n",
    "            *Tools.block(256, 512),\n",
    "            *Tools.block(512, 1024),\n",
    "            nn.Linear(1024, int(np.prod(opt.img_shape))),\n",
    "            nn.Tanh())\n",
    "        \n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr, betas=opt.b)\n",
    "        \n",
    "    def forward(self, noise, labels):\n",
    "        gen_input = torch.cat((self.label_emb(labels), noise), -1)\n",
    "        img = self.model(gen_input)\n",
    "        img = img.view(img.size(0), *opt.img_shape)\n",
    "        return img\n",
    "        \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9724604b-1ac2-4ceb-a049-bb826b26aec3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:33:00.801366Z",
     "iopub.status.busy": "2022-05-31T02:33:00.801197Z",
     "iopub.status.idle": "2022-05-31T02:33:00.807880Z",
     "shell.execute_reply": "2022-05-31T02:33:00.807344Z",
     "shell.execute_reply.started": "2022-05-31T02:33:00.801347Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.label_embedding = nn.Embedding(opt.classes, opt.classes)\n",
    "\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(opt.classes + int(np.prod(opt.img_shape)), 512),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Linear(512, 512),\n",
    "            nn.Dropout(0.4),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            nn.Linear(512, 1),\n",
    "        )\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr, betas=opt.b)\n",
    "        \n",
    "\n",
    "    def forward(self, img, labels):\n",
    "        \n",
    "        d_in = torch.cat((img.view(img.size(0), -1), self.label_embedding(labels)), -1)\n",
    "        validity = self.model(d_in)\n",
    "        return validity\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d3f43013-2776-4506-8ada-ea7b2674dc64",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:33:00.809056Z",
     "iopub.status.busy": "2022-05-31T02:33:00.808880Z",
     "iopub.status.idle": "2022-05-31T02:33:00.824658Z",
     "shell.execute_reply": "2022-05-31T02:33:00.824031Z",
     "shell.execute_reply.started": "2022-05-31T02:33:00.809037Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Dev:\n",
    "    def __init__(self):\n",
    "        self.g = Generator()\n",
    "        self.d = Discriminator()\n",
    "        self.loss = nn.MSELoss()\n",
    "        \n",
    "        self.models = (self.g, self.d)\n",
    "        \n",
    "        self.losses1 = []\n",
    "        self.losses2 = []\n",
    "        \n",
    "    def load_data(self):\n",
    "\n",
    "        def pretreat():\n",
    "            train = pd.read_csv('../../data/mnist/train.csv', header=None).values\n",
    "            # train, train_label = torch.tensor(train[...,1:]/255, dtype=torch.float), torch.tensor(train[...,0], dtype=torch.long)\n",
    "            train, train_label = torch.tensor((train[...,1:]-127.5)/127.5, dtype=torch.float), torch.tensor(train[...,0], dtype=torch.long)\n",
    "            return [train, train_label]\n",
    "        self.data = pretreat()\n",
    "        self.dataloader = torch.utils.data.DataLoader(\n",
    "            Data.TensorDataset(*self.data),\n",
    "            batch_size=opt.batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=opt.cpu_nums,\n",
    "            drop_last=True)\n",
    "        \n",
    "    def to_cuda(self):\n",
    "        opt.cuda = True\n",
    "        opt.device = opt.want_cuda\n",
    "        for i in self.models:\n",
    "            i.cuda()\n",
    "        if opt.big_cuda:\n",
    "            self.dataloader = torch.utils.data.DataLoader(\n",
    "                Data.TensorDataset(*[i.to(opt.device) for i in self.data]),\n",
    "                batch_size=opt.batch_size,\n",
    "                shuffle=True,\n",
    "                # num_workers=opt.cpu_nums,\n",
    "                drop_last=True)\n",
    "       \n",
    "    def to_cpu(self):\n",
    "        opt.cuda = False\n",
    "        opt.device = 'cpu'\n",
    "        for i in self.models:\n",
    "            i.cpu()\n",
    "        self.dataloader = torch.utils.data.DataLoader(\n",
    "            Data.TensorDataset(*self.data),\n",
    "            batch_size=opt.batch_size,\n",
    "            shuffle=True,\n",
    "            num_workers=opt.cpu_nums,\n",
    "            drop_last=True)\n",
    "        \n",
    "        \n",
    "    def train(self):\n",
    "        \n",
    "        losses1 = []\n",
    "        losses2 = []\n",
    "        \n",
    "        for i, (imgs, labels) in enumerate(self.dataloader):\n",
    "            \n",
    "            batch_size = imgs.shape[0]\n",
    "            valid = torch.ones(batch_size, 1, device=opt.device)\n",
    "            fake = torch.zeros(batch_size, 1, device=opt.device)\n",
    "            \n",
    "            self.g.optimizer.zero_grad()\n",
    "            z = torch.randn(batch_size, opt.latent_dim, device=opt.device)\n",
    "            gen_labels = torch.randint(0, opt.classes, (batch_size,), device=opt.device)\n",
    "            \n",
    "            gen_imgs = self.g(z, gen_labels)\n",
    "            validity = self.d(gen_imgs, gen_labels)\n",
    "            g_loss = self.loss(validity, valid)\n",
    "            g_loss.backward()\n",
    "            self.g.optimizer.step()\n",
    "        \n",
    "            self.d.optimizer.zero_grad()\n",
    "            validity_real = self.d(imgs, labels)\n",
    "            d_real_loss = self.loss(validity_real, valid)\n",
    "\n",
    "            validity_fake = self.d(gen_imgs.detach(), gen_labels)\n",
    "            d_fake_loss = self.loss(validity_fake, fake)\n",
    "\n",
    "            d_loss = (d_real_loss + d_fake_loss) / 2\n",
    "            d_loss.backward()\n",
    "            self.d.optimizer.step()\n",
    "            \n",
    "            losses1.append(g_loss.item())\n",
    "            losses2.append(d_loss.item())\n",
    "            \n",
    "        self.losses1.append(sum(losses1)/len(losses1))\n",
    "        self.losses2.append(sum(losses2)/len(losses2))\n",
    "        \n",
    "        print(self.losses1[-1], end=' ')\n",
    "        print(self.losses2[-1], end=' ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "550e8797-867c-4db5-8628-9ca63d19d811",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:33:00.985914Z",
     "iopub.status.busy": "2022-05-31T02:33:00.985629Z",
     "iopub.status.idle": "2022-05-31T02:33:04.125280Z",
     "shell.execute_reply": "2022-05-31T02:33:04.124630Z",
     "shell.execute_reply.started": "2022-05-31T02:33:00.985889Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "dev = Dev()\n",
    "dev.load_data()\n",
    "# dev.to_cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c6a63a06-bc71-4791-99b6-d888511378f3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-05-31T02:33:09.023937Z",
     "iopub.status.busy": "2022-05-31T02:33:09.023375Z",
     "iopub.status.idle": "2022-05-31T02:36:54.746894Z",
     "shell.execute_reply": "2022-05-31T02:36:54.746170Z",
     "shell.execute_reply.started": "2022-05-31T02:33:09.023908Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1 0.5961221016179292 0.13719373324121803 耗时11.170854091644287s\n",
      "epoch 2 0.5359262083935686 0.14323080359743193 耗时11.237585306167603s\n",
      "epoch 3 0.6133088479694139 0.11821600167542441 耗时11.079909324645996s\n",
      "epoch 4 0.6480698742322687 0.10453037231460087 耗时11.330132722854614s\n",
      "epoch 5 0.6597657604381825 0.10610322939216071 耗时11.376296520233154s\n",
      "epoch 6 0.6517164810982525 0.1076386373840336 耗时11.434344053268433s\n",
      "epoch 7 0.6378891641258174 0.11128970802339733 耗时11.264440059661865s\n",
      "epoch 8 0.645541713692439 0.1123869847983886 耗时11.376632452011108s\n",
      "epoch 9 0.6614856984880235 0.10835004317709523 耗时11.162298917770386s\n",
      "epoch 10 0.6784296927289066 0.09997127294285685 耗时11.222132921218872s\n",
      "epoch 11 0.6709121485423838 0.10483129679137825 耗时11.012072563171387s\n",
      "epoch 12 0.6897916725048652 0.1008377105761797 耗时11.169281959533691s\n",
      "epoch 13 0.652095741727668 0.1113354414383061 耗时11.425838708877563s\n",
      "epoch 14 0.6558075335481738 0.11250857356139737 耗时11.496849775314331s\n",
      "epoch 15 0.6438840158984193 0.11445540756496608 耗时11.37925934791565s\n",
      "epoch 16 0.637530365656329 0.12008023704601149 耗时11.386048078536987s\n",
      "epoch 17 0.6167932097983156 0.1231515149339142 耗时11.09522032737732s\n",
      "epoch 18 0.6146744389691924 0.1253534499078225 耗时11.316214323043823s\n",
      "epoch 19 0.6302509679116755 0.12019476775302847 耗时11.326010704040527s\n",
      "epoch 20 0.6087428892397473 0.12931499521956485 耗时11.431272983551025s\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "time0 = time.time()\n",
    "for epoch in range(20):\n",
    "    print(f'epoch {epoch+1}', end=' ')\n",
    "    dev.train()\n",
    "    time0, time_spend = time.time(), time.time()-time0\n",
    "    print(f'耗时{time_spend}s')\n",
    "    if not epoch % 1:\n",
    "        Tools.sample_image(10, epoch+1, dev.g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef07734d-ef95-45d9-b84b-ea6bd0d6e9ba",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
